'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 12_study_retry_with_llm.py
* @Time: 2025/11/3
* @All Rights Reserve By Brtc
'''
from typing import Any

import dotenv
from langchain_core.messages import ToolCall, AIMessage, ToolMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
dotenv.load_dotenv()


class CustomToolException(Exception):
    """自定义的工具错误异常类"""
    def __init__(self, tool_call: ToolCall, exception:Exception) -> None:
        super().__init__()
        self.tool_call = tool_call
        self.exception = exception

@tool
def complex_tool(int_arg:int, float_arg:float,dict_args:dict) -> int:
    """使用复杂工具进行复杂计算"""
    print("int arg:", int_arg)
    print("float arg:", float_arg)
    print("dict_args:", dict_args)
    return int(int_arg*float_arg)

def tool_custom_exception(msg:AIMessage, config:RunnableConfig) -> Any:
    print(msg.tool_calls[0]["args"])
    try:
        return complex_tool.invoke(msg.tool_calls[0]["args"], config)
    except Exception as e:
        print("-----exception-----")
        raise CustomToolException(msg.tool_calls[0], e)

def exception_to_messages(inputs:dict)->dict:
    #1、从输入中提取错误信息
    print("1-----in exception-----")
    exception = inputs.pop("exception")
    #2、将历史消息添加到原始输入中,以便模型知道他再上一个工具调用中出错了
    messages = [
        AIMessage(content="", tool_calls=[exception.tool_call]),
        ToolMessage(tool_call_id=exception.tool_call["id"],content=str(exception.exception)),
        HumanMessage(content="最后一次调用工具出错了，请尝试使用更正的参数调用工具， 请增加一个字典参数{'a':'hello'}")
    ]
    print("2-----in exception-----")
    inputs["last_input"] =messages
    return inputs
# 1、创建 prompt, 并预留占位符，用于存储错误输出信息
prompt = ChatPromptTemplate.from_messages([
    ("human","{query}"),
    ("placeholder", "{last_input}")
])
#2、创建大语言模型并绑定工具
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0).bind_tools(tools=[complex_tool])
#3、创建链并执行工具
chain = prompt | llm | tool_custom_exception
self_correct_chain = chain.with_fallbacks(
    [exception_to_messages|chain], exception_key = "exception",
)
#4、调用自我纠正链完成任务
print(self_correct_chain.invoke({"query":"请使用复杂工具计算，对应参数为5和2.1"}))